# -- coding:utf-8 --
import os
import pygame
from pygame.locals import *
from sys import exit

from collections import namedtuple
from copy import deepcopy

from uct import MCTS, Draw

N = 0
Z = 1
P = 2
M = 3
C = 4
X = 5
S = 6
J = 7
L = 8


R = 0
B = 1
A = 2

E = [1,0]

DOWN = R

SCORE_A = [0, 20, 55, 50, 100, 0, 0, 0]
SCORE_D = [0, 20, 55, 50, 100, 15, 20, 25]


class Utils:
    @staticmethod
    def pos_add(pos1, pos2):
        return (pos1[0] + pos2[0], pos1[1] + pos2[1])

    @staticmethod
    def pos_sub(pos1, pos2):
        return (pos1[0] - pos2[0], pos1[1] - pos2[1])

    @staticmethod
    def pos_flip(pos):
        return (pos[0], 9 - pos[1])

    @staticmethod
    def to_xy(pos):
        x = (pos[0]) * 60 + 14
        y = (pos[1]) * 60 + 16
        return (x, y)

    @staticmethod
    def to_pos(xy):
        px = (xy[0] - 14)//60
        py = (xy[1] - 16)//60
        return (px, py)

class Chess:
    State = namedtuple('State', ['board', 'winner', 'current_player', 'depth'])
    Step = namedtuple('Step', ['player', 'from_pos', 'to_pos'])

    @classmethod
    def initial_rule(cls):
        cls.__chess_man_info = []
        for player in range(A):
            object = []
            object.append(None)
            object.append(Chessman(player, Z))
            object.append(Chessman(player, P))
            object.append(Chessman(player, M))
            object.append(Chessman(player, C))
            object.append(Chessman(player, X))
            object.append(Chessman(player, S))
            object.append(Chessman(player, J))
            cls.__chess_man_info.append(object)

    @classmethod
    def initial_state(cls):
        board = []
        board.append(cls.make_board())
        board.append(cls.make_board())

        winner = None
        current_player = R

        return cls.State(board=board, 
            winner=winner, current_player=current_player, depth=0)

    @classmethod
    def apply_move(cls, state, move):
        board = deepcopy(state.board)
        winner = None

        _, type = cls.get_p_type(state, move.player, move.from_pos)
        _, type_e = cls.get_p_type(state, move.player, move.to_pos)

        if type_e == J:
            winner = state.current_player

        cls.set_p_type(board, move.player, move.from_pos, N)
        cls.set_p_type(board, move.player, move.to_pos, type)

        current_player = E[state.current_player]

        new_state = cls.State(board=board, 
            winner=winner, current_player=current_player, depth=state.depth+1)

        
        # print 'move', state.depth, move, new_state.winner, state.winner
        # print 'board'
        # for row in state.board[move.player]:
            # print row
        # print '=>'
        # for row in board[move.player]:
            # print row

           
        # error debug
        # if cls.get_error(new_state):
            # print 'error'
            # print 'old'
            # for row in state.board[move.player]:
                # print row
            # print 'new'
            # for row in board[move.player]:
                # print row
            # print 'old 2'
            # for row in state.board[E[move.player]]:
                # print row
            # print 'new 2'
            # for row in board[E[move.player]]:
                # print row                
            # print 'move', move
            # print 'type', type, type_e
            # print 'winner', winner
            # print 'depth', state.depth
            # exit()
            
        
        return new_state

    @classmethod
    def get_moves(cls, state):
        steps = []
        board = state.board[state.current_player]

        for py in range(10):
            for px in range(9):
                player, type = cls.get_p_type(state, state.current_player, (px, py))
                if player == state.current_player:
                    chessman = cls.__chess_man_info[state.current_player][type]
                    choices = chessman.choices(board, (px, py))
                    steps.extend( [ cls.make_step(state, (px, py), c) for c in choices  ] )
        
        # error debug
        if len(steps) == 0:
            print 'get_moves'
            for row in board:
                print row
            for py in range(10):
                for px in range(9):
                    player, type = cls.get_p_type(state, state.current_player, (px, py))
                    if player == state.current_player:
                        chessman = cls.__chess_man_info[state.current_player][type]
                        choices = chessman.choices(board, (px, py))  
                        print choices
            exit()    
                    
        return False, steps

    @classmethod
    def get_winner(cls, state, depth):
        if state.winner != None:
            return state.winner

        if depth > 150:
            return Draw

        if depth > 10:
            score_a, score_d = cls.get_score(state)
            if score_a[R] < 100 and score_a[B] < 100:
                return Draw
            elif score_a[R] - 150 > score_a[B]:
                return R
            elif score_a[B] - 150 > score_a[R]:
                return B


        return None

    @staticmethod
    def current_player(state):
        return state.current_player

    @classmethod
    def get_error(cls, state):
        ret = 0
        for player in range(A):
            for py in range(10):
                for px in range(9):
                    _, type = cls.get_p_type(state, player, (px, py))
                    if type == J or type == -J:
                        ret += 1
        
        return (ret < 4) and (state.winner == None)
            
        
    @classmethod
    def get_score(cls, state):
        board = state.board[state.current_player]
        
        score_a = [0,0]
        score_d = [0,0]
        for py in range(10):
            for px in range(9):
                player, type = cls.get_p_type(state, state.current_player, (px, py))
                score_a[player] += SCORE_A[type]
                score_d[player] += SCORE_D[type]

        return score_a, score_d      
                
        
    @staticmethod
    def make_board():
        list = [(Z, (0,3)), (Z, (2,3)), (Z, (4,3)), (Z, (6,3)), (Z, (8,3)), (P, (1,2)), (P, (7,2)),                   
            (C, (0,0)), (M, (1,0)), (X, (2,0)), (S, (3,0)), (J, (4,0)), (S, (5,0)), (X, (6,0)), (M, (7,0)), (C, (8,0))]

        board = []
        for _ in range(10):
            row = []
            for x in range(9):
                row.append(0)
            board.append(row)

        for l in list:
            type = l[0]
            pos = l[1]
            board[pos[1]][pos[0]] = type

        for l in list:
            type = l[0]
            pos = Utils.pos_flip(l[1])
            board[pos[1]][pos[0]] = -type
        return board

    @staticmethod
    def get_p_type(state, p, pos):
        px, py = pos
        type = state.board[p][py][px]
        if type > 0:
            return p, type
        else:
            return E[p], -type
        return E, N

    @staticmethod
    def set_p_type(board, player, pos, type):
        px, py = pos
        px_f, py_f = Utils.pos_flip(pos)
        
        board[player][py][px] = type
        board[E[player]][py_f][px_f] = -type

    @classmethod
    def make_step(cls, state, from_pos, to_pos):
        return cls.Step(player=state.current_player, from_pos=from_pos, to_pos=to_pos)

    @classmethod
    def apply_move_rule(cls, state, move):
        pl, type = cls.get_p_type(state, move.player, move.from_pos)

        if type != N and pl == move.player:
            chessman = cls.__chess_man_info[move.player][type]
            if chessman.is_in_rule(state.board[move.player], move.from_pos, move.to_pos):
                return cls.apply_move(state, move)
        return None

    @staticmethod
    def get_type(state, pos):
        px, py = pos
        type = state.board[R][py][px]
        if type > 0:
            return R, type
        else:
            return B, -type  
            
        return -1, N

    @classmethod
    def draw(cls, state):
        for py in range(10):
            for px in range(9):
                pos = (px, py)

                player, type = cls.get_type(state, pos)
                if type != N:
                    man = cls.__chess_man_info[player][type]
                    if DOWN == R:
                        pos = Utils.pos_flip(pos)
                    man.draw(pos)

class cplayer:
    def __init__(self, player, ai = None):
        self.player = player

        self.ai = ai
        self.ai_from_pos = None
                
    def do_ai(self, state, pos):
        if self.ai:
            return self.ai(state)
        elif pos:
            player, type = Chess.get_p_type(state, self.player, pos)
            if player == self.player and type != N:
                self.ai_from_pos = pos
            elif self.ai_from_pos:
                player, type = Chess.get_p_type(state, self.player, pos)
                if type == N or player != self.player:
                    return Chess.make_step(state, self.ai_from_pos, pos), None
        return None, None

class Chessman:
    type_flags = ['', 'Z', 'P', 'M', 'C', 'X', 'S', 'J' ]
    player_flags = ['R', 'B']
    
    def __init__(self, player, type):
        self.p = player
        self.image = pygame.image.load(os.path.join('res', Chessman.player_flags[player] + Chessman.type_flags[type] +'.png')).convert_alpha()
        self.flip = (player == DOWN)
        self.type = type
    
    def draw(self, pos):
        xy = Utils.to_xy(pos)
        SCREEN.blit(self.image, xy)

    def make_choices(self, board, pos, steps, rect, skips=None):
        ret = []

        for i in range(len(steps)):
            step = steps[i]
            skip = None
            if skips:
                skip = skips[i]

            to_pos = Utils.pos_add(pos, step)

            if rect[2] < to_pos[0] or rect[0] > to_pos[0] or rect[3] < to_pos[1] or rect[1] > to_pos[1]:
                continue
                
            if board[to_pos[1]][to_pos[0]] <= 0:
                if skip:
                    skip = Utils.pos_add(pos, skip)
                    if board[skip[1]][skip[0]] != 0:
                        continue

                ret.append(to_pos)

        return ret

    def make_choices_pc(self, board, pos, jump):
        ret = []
        rect = (0,0,8,9)
        
        loop = ( 
            (False,  1, 10, 1),
            (False, -1, -10, -1),
            (True,  1,   9, 1),
            (True,  -1, -9, -1),
            )
        
        for l in loop:
            status = False
            for i in range(l[1],l[2],l[3]):

                if l[0]:
                    step = (i, 0)
                else:
                    step = (0, i)

                to_pos = Utils.pos_add(pos, step)
                if rect[2] < to_pos[0] or rect[0] > to_pos[0] or rect[3] < to_pos[1] or rect[1] > to_pos[1]:
                    break
                
                type = board[to_pos[1]][to_pos[0]]

                if type == N and not status:
                    ret.append(to_pos)
                elif type < 0:
                    if status:
                        ret.append(to_pos)
                        break
                    if jump:
                        status = True
                    else:
                        ret.append(to_pos)
                        break
                elif type > 0:
                    if jump and not status:
                        status = True
                    else:
                        break
        return ret
        
    def choices(self, board, pos):
        ret = []
        if self.type == J:
            steps = ((0,1), (0,-1), (1,0), (-1,0))
            ret.extend(self.make_choices(board, pos, steps, (3,0,5,2)))
        elif self.type == S:
            steps = ((1,1), (1,-1), (-1,1), (-1,-1))
            ret.extend(self.make_choices(board, pos, steps, (3,0,5,2)))
        elif self.type == X:
            steps = ((2,2), (2,-2), (-2,2), (-2,-2))
            skips = ((1,1), (1,-1), (-1,1), (-1,-1))
            ret.extend(self.make_choices(board, pos, steps, (0,0,8,4), skips))
        elif self.type == M:
            steps = ((1,2), (-1,2),(1,-2), (-1,-2), (2,1), (2,-1),(-2,1), (-2,-1))
            skips = ((0,1), (0,1), (0,-1), (0,-1),  (1,0), (1,0), (-1,0), (-1,0))
            ret.extend(self.make_choices(board, pos, steps, (0,0,8,9), skips))
        elif self.type == Z:
            py = pos[1]
            if py > 4:
                steps = ((1,0), (-1,0), (0,1))
            else:
                steps = ((0,1), )
            ret.extend(self.make_choices(board, pos, steps, (0,0,8,9)))
        elif self.type == C:
            ret.extend(self.make_choices_pc(board, pos, False))
        elif self.type == P:
            ret.extend(self.make_choices_pc(board, pos, True))           
        return ret

    def is_in_rule(self, board, pos, step):
        choices = self.choices(board, pos)
        if step in choices:
            return True
        return False

def ai_A(state):
    result = ( MCTS(Chess, state).get_simulation_result(20))
    return result.move, result.root

def ai_B(state):
    result = ( MCTS(Chess, state).get_simulation_result(100))
    return result.move, result.root
    
pygame.init()

WIDTH = 560
HEIGHT = 620
PWIDTH = 200
BLACK = (0, 0, 0)
WHITE = (255,255,255)
FONT = pygame.font.SysFont("arial", 24)
SCREEN = pygame.display.set_mode((WIDTH+PWIDTH, HEIGHT), 0, 32)


REPORT = [0] * A

def update_message(state, note=None):
    pygame.draw.rect(SCREEN, BLACK, Rect((WIDTH,0), (PWIDTH,HEIGHT)))
    msg = []
    msg.append(FONT.render(u"Current Round: " + str(state.depth), True, WHITE))
    msg.append(FONT.render(u"Current Player: " + str(state.current_player), True, WHITE))
    if note:
        msg.append(None)
        msg.append(FONT.render(u"Red Player: " + str(int(note.win_rate(R)* 100)) + "%", True, WHITE))
        msg.append(FONT.render(u"Black Player: " + str(int(note.win_rate(B)* 100)) + "%", True, WHITE))
    msg.append(None)
    msg.append(FONT.render(u"Red Player: " + str(REPORT[R]), True, WHITE))
    msg.append(FONT.render(u"Black Player: " + str(REPORT[B]), True, WHITE))
    for y in range(len(msg)):
        if msg[y]:
            SCREEN.blit(msg[y], (WIDTH+20,0+30*y))
    
def main():
    
    pygame.display.set_caption("Chinese Chess")

    background = pygame.image.load(os.path.join('res', 'CP.jpg')).convert()

    # clock = pygame.time.Clock()

    players = []
    players.append(cplayer(R, None))
    players.append(cplayer(B, ai_A))

    Chess.initial_rule()

    root = None
    while True:
        state = Chess.initial_state()
        # print 'Chess Report:', REPORT

        while True:
            SCREEN.blit(background, (0,0))
            # x, y = pygame.mouse.get_pos()
            Chess.draw(state)
            
            update_message(state, root)

            pygame.display.update()

            input = None
            if state.winner != None or state.depth > 100:
                if state.winner != Draw and state.winner != None:
                    REPORT[state.winner] += 1
                break

            for event in pygame.event.get():
                if event.type == QUIT:
                    exit()
                elif event.type == MOUSEBUTTONDOWN:
                    mouse_down = event.button
                    mouse_down_x,mouse_down_y = event.pos
                    input = Utils.to_pos(event.pos)
                    if state.current_player == DOWN:
                        input = Utils.pos_flip(input)

            move, note = players[state.current_player].do_ai(state, input)
            if note:
                root = note
            if move != None:
                status = Chess.apply_move_rule(state, move)
                if status:
                    state = status
                    moves = Chess.get_moves(state)[1]



main()